import cv2
import numpy as np
from scipy.ndimage import gaussian_filter

class RetinalVesselSegmenter:
    def __init__(self, scales=[4.5, 3.2, 2.3], clahe_clip=2.0, clahe_grid=(8, 8)):
        """
        初始化参数
        :param scales: 多尺度高斯滤波的σ值列表
        :param clahe_clip: CLAHE对比度限制阈值
        :param clahe_grid: CLAHE分块网格大小
        """
        self.scales = scales
        self.clahe = cv2.createCLAHE(clipLimit=clahe_clip, tileGridSize=clahe_grid)
        self.alpha = 0.2  # 掩膜融合透明度

    def __call__(self, image):
        """
        执行血管分割并返回带掩膜的图像
        :param image: 输入图像，仅支持 numpy.ndarray 格式（BGR或RGB格式）
        :return: 带血管分割掩膜的 numpy 数组（BGR格式）
        """
        # 检查输入类型
        if not isinstance(image, np.ndarray):
            raise TypeError("输入必须是 numpy.ndarray 格式")
        if len(image.shape) != 3 or image.shape[2] != 3:
            raise ValueError("输入图像必须是3通道彩色图像")
        if image.dtype != np.uint8:
            image = (image * 255).astype(np.uint8)

        # 假设输入是RGB格式，转换为BGR
        if image.shape[2] == 3 and image[:, :, 0].max() <= 1:
            image = (image * 255).astype(np.uint8)
        bgr_image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

        # 处理流程
        enhanced = self.preprocess(bgr_image)
        vessel_mask = self.detect_vessels(enhanced)
        overlay = self.add_mask_to_image(bgr_image, vessel_mask)
        return overlay

    def preprocess(self, image):
        """
        图像预处理（CLAHE增强并提取G通道）
        :param image: 输入BGR图像
        :return: 预处理后的灰度图像
        """
        # 提取G通道并进行CLAHE增强
        g_channel = image[:, :, 1]
        enhanced = self.clahe.apply(g_channel)
        return enhanced

    def compute_hessian(self, image, sigma):
        """
        计算Hessian矩阵及其特征值
        :param image: 输入图像
        :param sigma: 高斯核尺度
        :return: 特征值图（λ1, λ2）、方向图
        """
        # 高斯二阶导数卷积
        gaussian = gaussian_filter(image, sigma)

        # 计算梯度
        grad_x = cv2.Sobel(gaussian, cv2.CV_64F, 1, 0, ksize=3)
        grad_y = cv2.Sobel(gaussian, cv2.CV_64F, 0, 1, ksize=3)

        # 计算二阶导数
        grad_xx = cv2.Sobel(grad_x, cv2.CV_64F, 1, 0, ksize=3)
        grad_xy = cv2.Sobel(grad_x, cv2.CV_64F, 0, 1, ksize=3)
        grad_yy = cv2.Sobel(grad_y, cv2.CV_64F, 0, 1, ksize=3)

        # 计算Hessian矩阵特征值
        lambda1 = np.zeros_like(image, dtype=np.float64)
        lambda2 = np.zeros_like(image, dtype=np.float64)
        direction = np.zeros_like(image, dtype=np.float64)

        for i in range(image.shape[0]):
            for j in range(image.shape[1]):
                H = np.array([[grad_xx[i, j], grad_xy[i, j]],
                              [grad_xy[i, j], grad_yy[i, j]]])
                eigvals = np.linalg.eigvalsh(H)
                lambda1[i, j] = np.abs(eigvals[0])
                lambda2[i, j] = np.abs(eigvals[1])
                direction[i, j] = np.arctan2(grad_y[i, j], grad_x[i, j])

        return lambda1, lambda2, direction

    def detect_vessels(self, enhanced_image):
        """
        多尺度血管检测主函数
        :param enhanced_image: 预处理后的图像
        :return: 血管分割结果（二值图）
        """
        vessel_masks = []

        for sigma in self.scales:
            lambda1, lambda2, direction = self.compute_hessian(enhanced_image, sigma)
            response = np.where(lambda2 > 1.2, lambda2, 0)  # Th=1.2
            response = self.non_max_suppression(response, direction)
            _, binary = cv2.threshold(response.astype(np.uint8), 0, 255, cv2.THRESH_BINARY)
            vessel_masks.append(binary)

        final_mask = self.fuse_masks(vessel_masks)
        return final_mask

    def non_max_suppression(self, response, direction):
        """
        非极大值抑制（方向自适应）
        :param response: 响应图
        :param direction: 方向图（弧度）
        :return: 抑制后的响应图
        """
        suppressed = np.zeros_like(response)
        for i in range(1, response.shape[0] - 1):
            for j in range(1, response.shape[1] - 1):
                theta = direction[i, j]
                sector = int(round(theta / (np.pi / 4))) % 4
                max_val = 0
                if sector == 0:
                    max_val = max(response[i, j - 1], response[i, j + 1])
                elif sector == 1:
                    max_val = max(response[i - 1, j + 1], response[i + 1, j - 1])
                elif sector == 2:
                    max_val = max(response[i - 1, j], response[i + 1, j])
                else:
                    max_val = max(response[i - 1, j - 1], response[i + 1, j + 1])

                if response[i, j] > max_val:
                    suppressed[i, j] = response[i, j]
        return suppressed

    def fuse_masks(self, masks):
        """
        多尺度掩模融合
        :param masks: 多尺度二值掩模列表
        :return: 融合后的最终掩模
        """
        fused = masks[0].copy()
        for mask in masks[1:]:
            fused = cv2.bitwise_or(fused, cv2.bitwise_and(mask, cv2.bitwise_not(fused)))
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
        fused = cv2.morphologyEx(fused, cv2.MORPH_CLOSE, kernel)
        return fused

    def add_mask_to_image(self, bgr_image, vessel_mask):
        """
        将二值掩膜与彩色图像叠加
        :param bgr_image: 输入的BGR图像
        :param vessel_mask: 血管二值掩膜
        :return: 叠加后的BGR图像
        """
        # 将掩膜转换为三通道绿色
        green_mask = np.zeros_like(bgr_image)
        green_mask[vessel_mask == 255] = [0, 255, 0]  # 绿色
        # 叠加图像
        overlay = cv2.addWeighted(bgr_image, 1, green_mask, self.alpha, 0)
        return overlay


# 测试
if __name__ == "__main__":
    segmenter = RetinalVesselSegmenter()

    # 示例：从文件加载图像并转换为numpy数组
    image = cv2.imread("test/19_left.jpg")  # 假设输入为BGR格式
    if image is None:
        raise ValueError("无法读取图像文件")

    result_image = segmenter(image)
    cv2.imwrite("segmented_image.jpg", result_image)